/*
 * tiles.cc
 *
 *  Created on: Mar 3, 2013
 *      Author: drewlewis
 */

#include <iostream>
#include <iomanip>
#include <chemistry/molecule/molecule.h>
#include <chemistry/qc/basis/integral.h>
#include <chemistry/qc/basis/split.h>
#include "tiled_array.h"
#include "Global.h"
#include <Eigen/Dense>
#include <ctime>

using namespace sc;
namespace TA = TiledArray;

void Get_Eigen_Hcore(const Ref<OneBodyInt> integral,
                 const Ref<GaussianBasisSet> basis, Matrix &H){
  const double *buffer = integral->buffer();
  // and compute core Hamiltonian integrals
  int nshell = basis->nshell();
  for(int s1=0; s1<nshell; s1++) {
    /// offset = first basis function in a given shell
    const int bf1_offset = basis->shell_to_function(s1);
    /// nbf1 gives total number of basis functions in a shell.
    const int nbf1 = basis->shell(s1).nfunction();

    for(int s2=0; s2<nshell; s2++) {
      const int bf2_offset = basis->shell_to_function(s2);
      const int nbf2 = basis->shell(s2).nfunction();

      integral->compute_shell(s1, s2);

      int bf12 = 0;
      for(int bf1=0; bf1<nbf1; ++bf1) {
        for(int bf2=0; bf2<nbf2; ++bf2, ++bf12) {
          H(bf1+bf1_offset, bf2+bf2_offset) = buffer[bf12];
        }
      }
    }
  }
}

// Import an eigen matrix with the one electron density into TA.
TA::Array<double, 2>::value_type
Eigen_to_TA(TA::Range const &range, const Matrix* Eigen_D){

  /// This is the tile which we are going to return.
  TA::Array<double, 2>::value_type tile(range);


  Eigen::Map<MatrixRM,Eigen::AutoAlign> map(tile.data(),
                           tile.range().size()[0],
                           tile.range().size()[1]);

  map = Eigen_D->block(tile.range().start()[0],
                       tile.range().start()[1],
                       tile.range().size()[0],
                       tile.range().size()[1]);

  return tile;
}

// Import an eigen matrix with the one electron density into TA.
TA::Array<double, 2>::value_type
MPQC_to_TA(TA::Range const &range, const Ref<SymmSCMatrix> &mpqc_mat){


    /// This is the tile which we are going to return.
    TA::Array<double, 2>::value_type tile(range);

    int t0start = tile.range().start()[0];
    int t1start = tile.range().start()[1];
    int t0size = tile.range().size()[0];
    int t1size = tile.range().size()[1];
    int t0end = t0start + t0size;
    int t1end = t1start + t1size;

    for(int i = t0start; i < t0end; ++i){

        int tile_i = i - t0start;
        for(int j = t1start; j < t1end; ++j){
            int tile_j = j - t1start;
            tile[tile_i * t1size + tile_j] = mpqc_mat->get_element(i,j);
        }
    }

    return tile;
}

void TA_to_Eigen(const TA::Array<double,2>::value_type &tile,
                 Matrix* Eigen_mat){

  Eigen_mat->block(tile.range().start()[0],
                   tile.range().start()[1],
                   tile.range().size()[0],
                   tile.range().size()[1]) =
              Eigen::Map<const MatrixRM>(tile.data(),
                                         tile.range().size()[0],
                                         tile.range().size()[1]);
}

void make_eigen_two_center(Ref<TwoBodyTwoCenterInt> Integral,
                           Matrix &e_mat){

  Ref<TwoBodyTwoCenterInt> integral_ = Integral->clone();
  const double *buffer = integral_->buffer();

  // and the two body two center integrals
  int nshell = Integral->basis1()->nshell();
  for(int s1=0; s1<nshell; s1++) {

    /// offset = first basis function in a given shell
    const int bf1_offset = Integral->basis1()->shell_to_function(s1);
    /// nbf1 gives total number of basis functions in a shell.
    const int nbf1 = Integral->basis1()->shell(s1).nfunction();

    for(int s2=0; s2<nshell; s2++) {
      const int bf2_offset = Integral->basis2()->shell_to_function(s2);
      const int nbf2 = Integral->basis2()->shell(s2).nfunction();

      integral_->compute_shell(s1, s2);

      int bf12 = 0;
      for(int bf1=0; bf1<nbf1; ++bf1) {
        for(int bf2=0; bf2<nbf2; ++bf2, ++bf12) {
          e_mat(bf1+bf1_offset, bf2+bf2_offset) = buffer[bf12];
        }
      }

    }
  }
}

void import_to_TA(madness::World &world, ARRAY2 &array,
                  const Matrix &eig_mat){

    ARRAY2::pmap_interface::const_iterator it = array.get_pmap()->begin();
    ARRAY2::pmap_interface::const_iterator end = array.get_pmap()->end();
    //const MatrixRM eig_row_mat = eig_mat;

    // Add local tiles
    for(; it != end; ++it) {

      // Construct a tile using a MADNESS task.
      madness::Future<ARRAY2::value_type> tile = world.taskq.add(
          &Eigen_to_TA, array.trange().make_tile_range(*it), &eig_mat);

      // Insert the tile into the array
      array.set(*it, tile);
    }
}


void import_to_TA(madness::World &world, ARRAY2 &array,
                  const Ref<SymmSCMatrix> &mpqc_mat){

    ARRAY2::pmap_interface::const_iterator it = array.get_pmap()->begin();
    ARRAY2::pmap_interface::const_iterator end = array.get_pmap()->end();

    // Add local tiles
    for(; it != end; ++it) {

      // Construct a tile using a MADNESS task.
      madness::Future<ARRAY2::value_type> tile = world.taskq.add(
          &MPQC_to_TA, array.trange().make_tile_range(*it), mpqc_mat);


      // Insert the tile into the array
      array.set(*it, tile);
    }
}



void export_from_TA(madness::World &world, const ARRAY2 &array,
                   Matrix &eigen_mat){

    ARRAY2::pmap_interface::const_iterator it = array.get_pmap()->begin();
    ARRAY2::pmap_interface::const_iterator end = array.get_pmap()->end();

    // Add local tiles
    for(; it != end; ++it) {
      const TA::Array<double,2>::value_type tile = array.Array::find(*it);

      // Construct a tile using a MADNESS task.
      world.taskq.add( &TA_to_Eigen, array.Array::find(*it), &eigen_mat);

    }
}



